version 17.0
clear all
cap log close
log using model_fit.log, replace

global RESULTS ../output
global TEMP ../temp
cap mkdir ../temp/estimate_w_xs
cap mkdir ../output/estimate_w_xs
cap mkdir ../output/estimate_w_xs/figures 

graph set window fontface default
graph set ps fontface default
graph set window fontfacemono "Consolas"
graph set ps fontfacemono "Consolas"
set scheme s1color


program main
	di "prep data"
	a_c_densities
	// moment_fit_by_regime
	model_fit_pooled_over_regimes	
	new_model_fit_by_regime
	cap log close
end

program a_c_densities
	cap qui import delimited $TEMP/estimate_w_xs/model_predict_disagg_no_xs.csv, encoding(ISO-8859-2) clear 
	keep a_i c_i
	tw kdensity a_i if inrange(a_i, -500000, 0), color(maroon) ///
		|| kdensity c_i if inrange(c_i, -500000, 0), color(navy) ///
      	subtitle("Density", position(11) size(medium)) xtitle("Value") ///
      	xscale(r(-500000 0)) xlab(#5) ///
      	legend(on order(- 1 "a_i" 2 "c_i") ///
      pos(10) size(small) ring(0) region(lp(solid)) cols(1)) plotregion(style(none)) scheme(s1color)
    graph export ../output/estimate_w_xs/figures/density_a_c_trim_new.pdf, replace

    tw kdensity a_i, color(maroon) ///
		|| kdensity c_i, color(navy) ///
      	subtitle("Density", position(11) size(medium)) xtitle("Value") xlab(#5) ///
      	legend(on order(- 1 "a_i" 2 "c_i") ///
      pos(10) size(small) ring(0) region(lp(solid)) cols(1)) plotregion(style(none)) scheme(s1color)
    graph export ../output/estimate_w_xs/figures/density_a_c_new.pdf, replace
end

program moment_fit_by_regime
	syntax
	* Prev model_predict_disagg_global
	cap qui import delimited $TEMP/estimate_w_xs/model_predict_disagg_no_xs.csv, encoding(ISO-8859-2) clear 
	
	gen sim_invY_post = pred_invasive if pred_nipt == 1 & wave == 3
	gen invY_post = did_invasive if did_nipt == 1 & wave == 3
	gen sim_invN_post = pred_invasive if pred_nipt == 0 & wave == 3
	gen invN_post = did_invasive if did_nipt == 0 & wave == 3
	gen nipt_post = did_nipt if wave == 3
	gen sim_nipt_post = pred_nipt if wave == 3
	gen sim_inv_w2 = pred_invasive if wave == 2

	rename bin_number bin50
	* 25 point bins
	local count = 1
	local bin_size = 25
	qui gen bin25 = .
	forval bin_max = `bin_size'(`bin_size')1000 {
			local bin_inf = `bin_max' - `bin_size'
			qui replace bin25 = `count' if fetus_risk <= `bin_max' & fetus_risk > `bin_inf'
			local count = `count' + 1
	}
	qui unique bin25
	local num_bins25 = r(unique)
	tempfile predict_clean
	save `predict_clean'
	
	foreach b_var in bin50 bin25 { 
		use `predict_clean', clear
		collapse (mean) sim_invY_post invY_post sim_invN_post invN_post sim_nipt_post nipt_post (count) pregnancy, ///
		  by(`b_var' policy_regime)
		di "`b_var'"
		if "`b_var'" == "bin50" replace `b_var' = `b_var' * 50
		if "`b_var'" == "bin25" replace `b_var' = `b_var' * 25
		rename pregnancy weight
		rename `b_var' bin_number 
		local reg_levs 1 3 4 5
		foreach rgm of local reg_levs {
			di `rgm'
			if `rgm' == 1 local rgm_name = "51_200"
			if `rgm' == 3 local rgm_name = "51_1000"
			if `rgm' == 4 local rgm_name = "hr_cov"
			if `rgm' == 5 local rgm_name = "1000"
			preserve 
			keep if policy_regime == `rgm'
			local ysc "0 1"
			local ylab "0(.1)1"
			foreach var in nipt_post invN_post invY_post {
				twoway (connect `var' bin_number, msymbol(d) mcolor(black)) ///
				(connect sim_`var' bin_number, msymbol(s) mcolor(red)), ///
				name(mfit_`b_var'_`var'_`rgm_name') ///
				xtitle("NT Screening result (estimated risk of chromosomal abnormalities)") ///
				xscale(reverse) xlabel(50 "1/50" 200 "1/200" 400 "1/400" 600 "1/600" ///
				800 "1/800" 1000 "1/1,000", format(%7.0fc)) ///
				yscale(r(`ysc')) ///
				ylabel(`ylab', format(%3.1fc)) ///
				legend(on order(1 "Data Moments" 2 "Model Moments") ///
				  pos(6) span col(2) row(1)) ///
				scale(0.9) /// 
				graphregion(color(white)) 
				graph export "${RESULTS}/estimate_w_xs/figures/mfit_`b_var'_`var'_`rgm_name'.pdf", ///
				  as(pdf) replace
			}

			gsort -bin_number
			qui export delimited $RESULTS/estimate_w_xs/mfit_`b_var'_`rgm_name'.csv, replace 
			restore 
		}
		
	}
end

program new_model_fit_by_regime
	cap qui import delimited $TEMP/estimate_w_xs/model_predict_disagg_no_xs.csv, encoding(ISO-8859-2) clear 
	
	rename bin_number bin50
	* 25 point bins
	local count = 1
	local bin_size = 25
	qui gen bin25 = .
	forval bin_max = `bin_size'(`bin_size')1000 {
			local bin_inf = `bin_max' - `bin_size'
			qui replace bin25 = `count' if fetus_risk <= `bin_max' & fetus_risk > `bin_inf'
			local count = `count' + 1
	}
	qui unique bin25
	local num_bins25 = r(unique)
	tempfile predict_clean
	
	gen inv_post = did_invasive if wave == 3
	gen inv_w2 = did_invasive if wave == 2
	gen sim_inv_post = pred_invasive if wave == 3
	gen sim_inv_w2 = pred_invasive if wave == 2
	gen nipt_post = did_nipt if wave == 3
	gen sim_nipt_post = pred_nipt if wave == 3
	save `predict_clean'
	
	// foreach b_var in bin50 bin25 { 
	foreach b_var in bin25 { 
		use `predict_clean', clear
		collapse (mean) inv_post sim_inv_post inv_w2 sim_inv_w2 nipt_post sim_nipt_post (count) pregnancy, by(`b_var' policy_regime)
		di "`b_var'"
		if "`b_var'" == "bin50" replace `b_var' = `b_var' * 50
		if "`b_var'" == "bin25" replace `b_var' = `b_var' * 25
		rename pregnancy weight
		rename `b_var' bin_number 
		local reg_levs 1 3 4 5
		foreach rgm of local reg_levs {
			di `rgm'
			if `rgm' == 1 local rgm_name = "51_200"
			if `rgm' == 3 local rgm_name = "51_1000"
			if `rgm' == 4 local rgm_name = "hr_cov"
			if `rgm' == 5 local rgm_name = "1000"
			preserve 
			keep if policy_regime == `rgm'
			local ysc "0 1"
			local ylab "0(.1)1"
			cap graph drop model_fit_inv_w2_xs
			cap graph drop model_fit_inv_post_xs
			cap graph drop model_fit_nipt_post_xs
			foreach var in inv_w2 inv_post nipt_post {
				twoway (connect `var' bin_number, msymbol(d) mcolor(black)) ///
				(connect sim_`var' bin_number, msymbol(s) mcolor(red)), ///
				name(new_`b_var'_`var'_`rgm_name') xtitle("NT Screening result (estimated risk of chromosomal abnormalities)") ///
				xscale(reverse) xlabel(50 "1/50" 200 "1/200" 400 "1/400" 600 "1/600" ///
				800 "1/800" 1000 "1/1,000", format(%7.0fc)) ///
				yscale(r(`ysc')) ///
				ylabel(`ylab', format(%3.1fc)) ///
				legend(on order(1 "Data Moments" 2 "Model Moments") ///
				  pos(6) span col(2) row(1)) ///
				scale(0.9) /// 
				graphregion(color(white)) 
				graph export "${RESULTS}/estimate_w_xs/figures/new_mfit_`var'_`rgm_name'.pdf", as(pdf) replace
			}
			gsort -bin_number
			qui export delimited $RESULTS/estimate_w_xs/new_`b_var'_`rgm_name'.csv, replace 
			restore 
	
		}
	} 
end 

program model_fit_pooled_over_regimes
	cap qui import delimited $TEMP/estimate_w_xs/model_predict_disagg_no_xs.csv, encoding(ISO-8859-2) clear 

	*qui import delimited $TEMP/estimate_w_xs/model_predict_disagg.csv, encoding(ISO-8859-2) clear 
	gen inv_post = did_invasive if wave == 3
	gen inv_w2 = did_invasive if wave == 2
	gen sim_inv_post = pred_invasive if wave == 3
	gen sim_inv_w2 = pred_invasive if wave == 2
	gen nipt_post = did_nipt if wave == 3
	gen sim_nipt_post = pred_nipt if wave == 3
	
	* 25 point bins
	local count = 1
	local bin_size = 25
	qui drop bin_number
	qui gen bin_number = .
	forval bin_max = `bin_size'(`bin_size')1000 {
			local bin_inf = `bin_max' - `bin_size'
			qui replace bin_number = `count' if fetus_risk <= `bin_max' & fetus_risk > `bin_inf'
			local count = `count' + 1
	}
	qui unique bin_number
	local num_bins = r(unique)
	
	* collapse data by 25 point bin
	collapse inv_post sim_inv_post inv_w2 sim_inv_w2 nipt_post sim_nipt_post (count) pregnancy, by(bin_number)
	qui replace bin_number = bin_number * 25
	
	* make plots
	cap graph drop model_fit_inv_w2
	local ysc "0 1"
	local ylab "0(.1)1"
	foreach var in inv_w2 inv_post nipt_post {
		twoway (connect `var' bin_number, msymbol(d) mcolor(black)) ///
		(connect sim_`var' bin_number, msymbol(s) mcolor(red)), ///
		name(model_fit_`var') xtitle("NT Screening result (estimated risk of chromosomal abnormalities)") ///
		xscale(reverse) xlabel(50 "1/50" 200 "1/200" 400 "1/400" 600 "1/600" ///
	    800 "1/800" 1000 "1/1,000", format(%7.0fc)) ///
		yscale(r(`ysc')) ///
		ylabel(`ylab', format(%3.1fc)) ///
		legend(on order(1 "Data Moments" 2 "Model Moments") ///
		  pos(6) span col(2) row(1)) ///
		scale(0.9) /// 
		graphregion(color(white)) 
		graph export "${RESULTS}/estimate_w_xs/figures/mfit_`var'_pooled.pdf", as(pdf) replace
	}

	gsort -bin_number
    qui export delimited $RESULTS/estimate_w_xs/pooled_moment_fit.csv, replace 
end

program prep_moment_data_testing
	gen error = data_moments - model_moments
	drop if strpos(name, "policy") != 0 
	qui gen bin_number = substr(name, -2, 2)
	qui replace bin_number = substr(name, -1, 1) if strpos(bin_number, "_") != 0
	destring bin_number, replace
	
	gen nipt = 1 * strpos(name, "nipt") != 0
	gen invY = 1 *  (v1 > 20 & v1 <= 40)
	gen invN = 1 * (v1 > 40 & v1 <= 80) 
	gen wave  = 2 if strpos(name, "w2") != 0
	replace wave = 3 if strpos(name, "w3") != 0
end


* Execute
main

